import warnings
warnings.filterwarnings('ignore')
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(40)

from tensorflow.keras import layers, models

tf.random.set_seed(777)

import os
import cv2
import numpy as np
import matplotlib.pyplot as plt
import random

char_set = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']

# data sets
IMG_HEIGHT = 60
IMG_WIDTH = 160
char_num = 4
characters = len(char_set)
"""
验证码标签为4*10维的矩阵 label: 1327
    [[0,1,0,0,0,0,0,0,0,0],
    [0,0,0,1,0,0,0,0,0,0],
    [0,0,1,0,0,0,0,0,0,0],
    [0,0,0,0,0,0,0,1,0,0]]
"""
def label2mat(label):
    label_mat = np.zeros((4,10))
    for i, num in enumerate(label):
        idx = int(num)
        label_mat[i][idx] = 1
    return label_mat

def readData(file_path):
    x_images = []
    y_labels = []
    for item in os.listdir(file_path):
        if item == '.nomedia':
            continue
        item_path = os.path.join(file_path, item)
        image = cv2.imread(item_path)/255
        x_images.append(image)
        label = os.path.splitext(item)[0]
        y_labels.append(label2mat(label))
    return np.array(x_images), np.array(y_labels)

test_dir = "../../../../../large_data/DL1/_many_files/vcode_data/test"
# test_dir = "../../../../../large_data/DL1/_many_files/vcode_data/train"
x_images, y_labels = readData(test_dir)

def mat2text(mat):
    text = []
    for i, c in enumerate(mat):
        text.append(char_set[c])
    return "".join(text)

plt.figure(figsize=(10, 5))
for i in range(20):
    plt.subplot(5, 4, i + 1)
    plt.xticks([])
    plt.yticks([])
    plt.grid(False)
    r = int(random.randint(0, x_images.shape[0] - 21))
    plt.imshow(x_images[i+r, :])
    plt.title(mat2text(np.argmax(y_labels[i+r:(i+r+1)], axis=2)[0]))
    plt.axis("off")
plt.show()

from sklearn.model_selection import train_test_split
x_train, x_test, y_train, y_test = train_test_split(x_images, y_labels, train_size=0.95)

# data pipeline
db_train = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(batch_size=16)
db_test = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(batch_size=16)

model = models.Sequential([
    layers.Conv2D(6, (3, 3), activation='relu', input_shape=(60, 160, 3)),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(16, (3, 3), activation='relu'),
    layers.MaxPooling2D((2, 2)),

    layers.Flatten(),
    layers.Dense(516, activation='relu'),

    layers.Dense(40),
    layers.Reshape([4, 10]),
    layers.Softmax()
])

model.summary()

model.compile(optimizer="adam",
              loss='categorical_crossentropy',
              metrics=['accuracy'])

epochs = 20
history = model.fit(
    db_train,
    validation_data=db_test,
    epochs=epochs)

# test a sample
r1 = int(random.randint(0, x_test.shape[0] - 1))
predictions = model.predict(x_test[r1:r1+1])

print("predictions:", mat2text(np.argmax(predictions, axis=2)[0]))
print("y_test:", mat2text(np.argmax(y_test[r1:r1+1], axis=2)[0]))

# 保存模型
os.makedirs('model', exist_ok=True)
model.save('model/12_model.h5')

# 六、模型评估
acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

loss = history.history['loss']
val_loss = history.history['val_loss']

epochs_range = range(epochs)

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)

plt.plot(epochs_range, acc, label='Training Accuracy')
plt.plot(epochs_range, val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, loss, label='Training Loss')
plt.plot(epochs_range, val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.show()

# del model
# 加载模型
new_model = tf.keras.models.load_model('model/12_model.h5')
new_model.summary()
loss, acc = new_model.evaluate(x_test, y_test)
print("Restored model, accuracy: {:5.2f}%".format(100*acc))

plt.figure(figsize=(10, 5))
for i in range(10):
    plt.subplot(5, 2, i + 1)
    plt.xticks([])
    plt.yticks([])
    plt.grid(False)
    r = int(random.randint(0, x_test.shape[0] - 11))
    plt.imshow(x_test[i+r, :])
    label_y = mat2text(np.argmax(y_test[i + r:(i + r + 1)], axis=2)[0])
    # predictions = model.predict(x_test[i + r:i + r + 1] / 255)
    predictions = new_model.predict(x_test[i + r:i + r + 1]) #加载模型
    predict_y = mat2text(np.argmax(predictions, axis=2)[0])
    print("i:", i+1, "label_y:", label_y, "predict_y:", predict_y)
    # 使用模型预测验证码
    plt.title(predict_y)
    plt.axis("off")
plt.show()
